#!/bin/bash
#SBATCH --gpus=2
#SBATCH -x paraai-n32-h-01-agent-[1-33],paraai-n32-h-01-agent-[48-56],paraai-n32-h-01-agent-[63-197]
export PYTHONUNBUFFERED=1

model=qwen1.8b.chat
new_date=$(date +%Y-%m-%d)_${model}
if [ ! -d "log/$new_date" ]; then
    mkdir -p "log/$new_date"
fi

model_name=Qwen-1_8B-Chat
data_type=CKnowEdit

module load compilers/cuda/12.1
module load cudnn/8.8.1.3_cuda12.x
module load compilers/gcc/11.3.0
source activate ke2torch23cu121
export HUGGINGFACE_CACHE=/home/bingxing2/home/scx7avs/lyc/huggingface/

batch_size=10
prob_type=target_new
layer_idx=0,24
max_tokens=128
max_sample_steps=512
data_name=type7_277

i=0
# ff_attrs=attn.c_attn
# mlp.c_proj,attn.c_proj,mlp.w1,mlp.w2,attn.c_attn
for ff_attrs in {mlp.c_proj,attn.c_proj}; do
    echo "$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps"
    CUDA_VISIBLE_DEVICES=$i python kn.py \
        --model_name $model_name \
        --batch_size $batch_size \
        --data_type $data_type \
        --data_path ../../EasyEdit/dataset/ccks2024_know_edit/CKnowEditType/$data_name.json \
        --transformer_layers_attr transformer.h \
        --ff_attrs $ff_attrs \
        --prob_type $prob_type \
        --layer_idx $layer_idx \
        --steps 20 \
        --max_tokens $max_tokens \
        --max_sample_steps $max_sample_steps \
        > log/$new_date/$i-$data_name-$model-$layer_idx-bs$batch_size-$prob_type-$ff_attrs-$max_tokens-$max_sample_steps-1.log 2>&1 &
        i=$((i+1))
done
wait